package org.example.util;

import jakarta.annotation.PostConstruct;
import org.apache.commons.lang3.ArrayUtils;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.example.processor.FeatureStorage;
import org.nd4j.common.io.ClassPathResource;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;
import org.nd4j.shade.guava.util.concurrent.ThreadFactoryBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Stream;
@Service
public class XrFileUtil {
    // 线程池配置
    private static final ExecutorService featureExecutor = Executors.newFixedThreadPool(
            6,
            new ThreadFactoryBuilder().setNameFormat("feature-pool-%d").build()
    );
    @Autowired
    private ResourceLoader resourceLoader;
    @Autowired
    private ComputationGraph model;

    @PostConstruct
    public void processDirectory() {
        try {
            int batchSize = 100;
            Resource resource = resourceLoader.getResource("classpath:img/"); // 替换你的文件夹名
            File folder = resource.getFile(); // 获取 File 对象
            System.out.println("资源目录路径: " + folder.getAbsolutePath());
            ClassPathResource classPathResource = new ClassPathResource(folder.getAbsolutePath());
            Stream<Path> paths = Files.walk(Paths.get(classPathResource.getPath()));
            List<CompletableFuture<Void>> futures = new ArrayList<>();

            paths.filter(Files::isRegularFile)
//                    .filter(p -> p.toString().endsWith(".jpg"))
                    .forEach(p -> {
                        CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                            try {
                                System.out.println(p.toAbsolutePath().toString());
                                processSingleImage(p, model);
                            } catch (Exception e) {
                                handleError(p, e);
                            }
                        }, featureExecutor);
                        futures.add(future);

                        // 批量提交控制
                        if (futures.size() >= batchSize) {
                            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
                            futures.clear();
                        }
                    });

            // 处理剩余任务
            if (!futures.isEmpty()) {
                CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
            }
        } catch (IOException e) {
            throw new RuntimeException("目录遍历失败", e);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void processSingleImage(Path imagePath, ComputationGraph model) throws Exception {
        // 特征提取
        INDArray feature = extractFeatures(imagePath, model);

        // 向量转换
//        List<Float> vector = convertToVector(feature);
        FeatureStorage.storeFeature(imagePath.toAbsolutePath().toString(), feature);
        // 数据库插入
//        insertToMilvus(
//                imagePath.getFileName().toString(),
//                vector,
//                parseMetadata(imagePath)
//        );
    }

    private INDArray extractFeatures(Path imagePath, ComputationGraph model) throws Exception {
        INDArray image = new NativeImageLoader(224, 224, 3)
                .asMatrix(imagePath.toFile());
        new VGG16ImagePreProcessor().transform(image);
        model.setInputs(image);
        return model.feedForward().get("fc2"); // 获取倒数第二层输出作为特征
    }

    private List<Float> convertToVector(INDArray ndArray) {
        float[] floatArray = ndArray.toFloatVector();
        return Arrays.asList(ArrayUtils.toObject(floatArray));
    }

    // 错误处理（可接入监控系统）
    private void handleError(Path path, Throwable e) {
        System.err.printf("处理失败: %s → %s%n", path, e.getMessage());
        // 记录失败文件路径，用于重试
    }
}
